import timm


# 根据参数切换模型
def create_model(name, in_channels, num_classes, use_timm=False):
    name = name.lower()
    if (use_timm):
        return timm.create_model(name, in_chans=in_channels, num_classes=10)
    if (name == 'lenet'):
        from models.LeNet import LeNet as Model
    elif (name == 'alexnet'):
        from models.AlexNet import AlexNet as Model
    elif (name == 'inceptionv1'):
        from models.InceptionV1 import InceptionV1 as Model
    elif (name == 'vgg11'):
        from models.VGG import VGG11 as Model
    elif (name == 'vgg13'):
        from models.VGG import VGG13 as Model
    elif (name == 'vgg16'):
        from models.VGG import VGG16 as Model
    elif (name == 'vgg' or name == 'vgg16c'):
        from models.VGG import VGG16C as Model
    elif (name == 'vgg19'):
        from models.VGG import VGG19 as Model
    elif (name == 'resnet18'):
        from models.ResNet import ResNet18 as Model
    elif (name == 'resnet34'):
        from models.ResNet import ResNet34 as Model
    elif (name == 'resnet' or name == 'resnet50'):
        from models.ResNet import ResNet50 as Model
    elif (name == 'resnet101'):
        from models.ResNet import ResNet101 as Model
    elif (name == 'resnet152'):
        from models.ResNet import ResNet152 as Model
    elif (name == 'densenet121'):
        from models.DenseNet import DenseNet121 as Model
    elif (name == 'densenet169'):
        from models.DenseNet import DenseNet169 as Model
    elif (name == 'densenet201'):
        from models.DenseNet import DenseNet201 as Model
    elif (name == 'densenet264'):
        from models.DenseNet import DenseNet264 as Model
    else:
        from models.SimpleNet import SimpleNet as Model
    return Model(num_classes, in_channels)